package ysgrpc

import (
	"context"
	"gitee.com/kmyss/gf-ex/ysgrpc/proto"
	"gitee.com/kmyss/gf-ex/yslog"
	"gitee.com/kmyss/gf-ex/ysother"
	"github.com/gogf/gf/v2/container/garray"
	"github.com/gogf/gf/v2/container/gmap"
	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/text/gstr"
	"google.golang.org/grpc"
	"sync"
)

func RegisterMqServer(registrar grpc.ServiceRegistrar) {
	ysother.ShowGRPCDes(proto.MQ_ServiceDesc)
	registrar.RegisterService(&proto.MQ_ServiceDesc, &MQServer{
		clients: gmap.NewIntAnyMap(true),
		topics:  gmap.NewStrAnyMap(true),
	})
}

type clientInfo struct {
	id     int
	s      proto.MQ_SubServer
	topics string
	cancel context.CancelFunc
}

type MQServer struct {
	l *glog.Logger
	proto.UnimplementedMQServer

	clients *gmap.IntAnyMap
	topics  *gmap.StrAnyMap
	nextId  int
	mutex   sync.Mutex
}

func (s *MQServer) log() *glog.Logger {
	if s.l == nil {
		s.l = yslog.Log(GenSubCfgNodeName("mq.server"))
	}
	return s.l
}

func (s *MQServer) Sub(request *proto.SubRequest, subServer proto.MQ_SubServer) error {
	topics := gstr.SplitAndTrim(request.GetTopics(), ",")

	// topics 为空时认为是订阅所有话题
	if len(topics) == 0 {
		topics = append(topics, "")
	}

	// 构建客户端对象
	// 这里需要使用锁
	s.mutex.Lock()
	ctx, cancel := context.WithCancel(subServer.Context())
	c := clientInfo{
		id:     s.nextId,
		s:      subServer,
		topics: request.GetTopics(),
		cancel: cancel,
	}
	s.log().Infof(ctx, "(%d)订阅话题: [ %s ]", c.id, request.GetTopics())
	s.clients.Set(c.id, &c)
	s.nextId++
	s.mutex.Unlock()

	// 发送客户端 id
	encode, err := gjson.Encode(g.MapStrInt{
		"id": c.id,
	})
	if err != nil {
		s.log().Error(ctx, "ID 封装异常", err)
		return err
	}
	err = subServer.Send(&proto.Message{
		Topic:   "ok.sub",
		Message: encode,
	})
	if err != nil {
		s.log().Error(ctx, "ID 发送异常", err)
		return err
	}

	// 注册 topic 与 客户端关联
	for _, topic := range topics {

		if !s.topics.Contains(topic) {
			s.topics.Set(topic, garray.NewIntArray(true))
		}
		s.topics.Get(topic).(*garray.IntArray).Append(c.id)
	}

	select {
	case <-ctx.Done():
	}

	// 把订阅的所有 topic 删除
	for _, topic := range topics {
		arr := s.topics.Get(topic).(*garray.IntArray)
		arr.RemoveValue(c.id)
		if arr.IsEmpty() {
			s.topics.Remove(topic)
		}

	}

	s.log().Infof(ctx, "(%d)取消订阅: [%s]", c.id, c.topics)

	return nil
}

func (s *MQServer) Pub(ctx context.Context, request *proto.Message) (*proto.PubReply, error) {
	num := 0
	for topic, servers := range s.topics.MapStrAny() {
		if gstr.HasPrefix(request.GetTopic(), topic) {
			for _, id := range servers.(*garray.IntArray).Slice() {
				err := s.clients.Get(id).(*clientInfo).s.Send(request)
				if err != nil {
					yslog.Error(err)
					continue
				}
				num++
			}
		}
	}

	return &proto.PubReply{Num: int32(num)}, nil
}

func (s *MQServer) UnSub(ctx context.Context, request *proto.UnSubRequest) (*proto.Message, error) {
	id := int(request.Id)

	if !s.clients.Contains(id) {
		s.log().Error(ctx, "取消的 ID 错误！", id)
		return nil, gerror.New("取消的 ID 异常！")
	}

	c := s.clients.Get(id).(*clientInfo)
	c.cancel()

	s.clients.Remove(id)

	return &proto.Message{
		Topic:   "ok.unsub",
		Message: nil,
	}, nil
}
